//! WebSocket 处理器模块
//! 实现基于 WebSocket 的实时通信

use axum::{
    extract::{ws::{Message, WebSocket}, WebSocketUpgrade, Query},
    response::Response,
    routing::get,
    Router,
};
use futures::{SinkExt, StreamExt};
use serde::Deserialize;
use std::sync::Arc;
use tokio::sync::{mpsc, RwLock, Mutex};
use uuid::Uuid;


use super::connection::{ConnectionManagerHandle, ConnectionType, WebSocketSender,ConnectionManager};
use super::message::{HubMessage, HubResponse, MessageSerializer, MessageType};
use crate::{
    models::token_model::TokenModel,
    services::RealtimeAuthMiddleware,
};

/// WebSocket 连接查询参数
#[derive(Debug, Deserialize)]
pub struct WebSocketQuery {
    pub access_token: Option<String>,
    pub token: Option<String>,
}

/// WebSocket 处理器
#[derive(Clone)]
pub struct WebSocketHandler {
    connection_manager: ConnectionManagerHandle,
    auth_middleware: RealtimeAuthMiddleware,
}

impl WebSocketHandler {
    pub fn new(connection_manager: ConnectionManagerHandle, auth_middleware: RealtimeAuthMiddleware) -> Self {
        Self {
            connection_manager,
            auth_middleware,
        }
    }
    
    /// 创建WebSocket路由
    pub fn create_route(&self) -> Router {
        let handler = self.clone();
        Router::new().route("/ws", get(move |ws: WebSocketUpgrade, query: Query<WebSocketQuery>| async move {
            ws.on_upgrade(move |socket| handler.handle_connection(socket, query))
        }))
    }
    
    /// 处理 WebSocket 连接（带身份验证）
    pub async fn handle_connection(self, socket: WebSocket, query: Query<WebSocketQuery>) {
        // 身份验证
        let token_model = match self.authenticate_connection(&query).await {
            Ok(token) => token,
            Err(e) => {
                tracing::error!("WebSocket 连接身份验证失败: {}", e);
                return;
            }
        };
        
        // 调用内部处理逻辑
        self.handle_websocket_internal(socket, token_model).await;
    }
    
    /// 身份验证连接
    async fn authenticate_connection(&self, query: &Query<WebSocketQuery>) -> anyhow::Result<TokenModel> {
        // 构建查询参数字符串用于身份验证
        let query_string = if let Some(token) = query.access_token.as_ref() {
            format!("access_token={}", token)
        } else if let Some(token) = query.token.as_ref() {
            format!("token={}", token)
        } else {
            return Err(anyhow::anyhow!("缺少访问令牌"));
        };
        
        // 使用身份验证中间件验证Token
        self.auth_middleware.authenticate_websocket(&query_string).await
    }
    
    /// 处理 WebSocket 连接（内部实现）
    async fn handle_websocket_internal(self, socket: WebSocket, token_model: TokenModel) {
        let (ws_sender, mut ws_receiver) = socket.split();
        
        // 使用 Arc 和 Mutex 共享 ws_sender
        let ws_sender = Arc::new(Mutex::new(ws_sender));
        
        // 创建消息通道
        let (tx, mut rx) = mpsc::unbounded_channel();
        
        // 创建连接发送器
        let connection_id = Uuid::new_v4();
        let sender = WebSocketSender {
            connection_id,
            sender: tx.clone(),
        };
        
        // 保存用户ID用于日志记录
        let user_id = token_model.user_id.clone();
        
        // 注册连接（包含身份验证信息）
        let connection_id = self.connection_manager
            .add_connection_with_token(sender, ConnectionType::WebSocket, Some(token_model))
            .await;
        
        tracing::info!("WebSocket 连接建立: {} (用户: {})", connection_id, user_id);
        
        // 发送连接确认消息
        let connect_response = HubResponse::success(
            None,
            serde_json::json!({
                "connectionId": connection_id.to_string(),
                "protocol": "signalr",
                "version": 1
            })
        );
        
        if let Ok(message) = MessageSerializer::serialize_response(&connect_response) {
            let mut sender = ws_sender.lock().await;
            if let Err(e) = sender.send(Message::Text(message)).await {
                tracing::error!("发送连接确认消息失败: {}", e);
                // 清理连接
                self.connection_manager.remove_connection(connection_id).await;
                return;
            }
            // 等待消息发送完成
            if let Err(e) = sender.flush().await {
                tracing::error!("刷新连接确认消息失败: {}", e);
                // 清理连接
                self.connection_manager.remove_connection(connection_id).await;
                return;
            }
        }
        
        // 处理接收和发送消息
        let connection_manager = self.connection_manager.clone();
        
        // 发送任务
        let send_task_ws_sender = ws_sender.clone();
        let send_task = tokio::spawn(async move {
            while let Some(response) = rx.recv().await {
                if let Ok(message_text) = MessageSerializer::serialize_response(&response) {
                    let mut sender = send_task_ws_sender.lock().await;
                    if sender.send(Message::Text(message_text)).await.is_err() {
                        break;
                    }
                }
            }
        });
        
        // 广播消息监听任务
        let broadcast_task_ws_sender = ws_sender.clone();
        let mut broadcast_receiver = connection_manager.subscribe();
        let broadcast_task = tokio::spawn(async move {
            while let Ok(response) = broadcast_receiver.recv().await {
                if let Ok(message_text) = MessageSerializer::serialize_response(&response) {
                    let mut sender = broadcast_task_ws_sender.lock().await;
                    if sender.send(Message::Text(message_text)).await.is_err() {
                        break;
                    }
                }
            }
        });
        
        // 接收任务
        let receive_task_ws_sender = ws_sender.clone();
        let receive_task = {
            let connection_manager = connection_manager.clone();
            
            tokio::spawn(async move {
                while let Some(result) = ws_receiver.next().await {
                    match result {
                        Ok(Message::Text(text)) => {
                            // 更新连接活动时间
                            if let Err(e) = connection_manager.update_connection_activity(connection_id).await {
                                tracing::error!("更新连接活动时间失败: {}", e);
                                break;
                            }
                            
                            // 解析消息
                            match MessageSerializer::parse_hub_message(&text) {
                                Ok(hub_message) => {
                                    // 处理消息
                                    if let Err(e) = Self::handle_message(
                                        &connection_manager, 
                                        connection_id, 
                                        hub_message
                                    ).await {
                                        tracing::error!("处理消息失败: {}", e);
                                        break;
                                    }
                                }
                                Err(e) => {
                                    tracing::error!("解析消息失败: {}", e);
                                    
                                    // 发送错误响应
                                    let error_response = HubResponse::error(
                                        None,
                                        format!("消息格式错误: {}", e)
                                    );
                                    
                                    if let Ok(message_text) = MessageSerializer::serialize_response(&error_response) {
                                        let _ = tx.send(HubResponse::error(None, message_text));
                                    }
                                }
                            }
                        }
                        Ok(Message::Close(_)) => {
                            tracing::info!("WebSocket 连接关闭: {}", connection_id);
                            // 立即清理连接
                            connection_manager.remove_connection(connection_id).await;
                            break;
                        }
                        Ok(Message::Ping(data)) => {
                            // 响应 Ping 消息
                            let mut sender = receive_task_ws_sender.lock().await;
                            if sender.send(Message::Pong(data)).await.is_err() {
                                break;
                            }
                        }
                        Ok(Message::Pong(_)) => {
                            // 忽略 Pong 消息
                        }
                        Ok(Message::Binary(_)) => {
                            tracing::warn!("收到二进制消息，暂不支持");
                        }
                        Err(e) => {
                            tracing::error!("WebSocket 接收错误: {}", e);
                            break;
                        }
                    }
                }
            })
        };
        
        // 等待任一任务完成
        tokio::select! {
            result = send_task => {
                if let Err(e) = result {
                    tracing::error!("发送任务错误: {}", e);
                }
            },
            result = receive_task => {
                if let Err(e) = result {
                    tracing::error!("接收任务错误: {}", e);
                }
            },
            result = broadcast_task => {
                if let Err(e) = result {
                    tracing::error!("广播任务错误: {}", e);
                }
            },
        }
        
        // 清理连接
        connection_manager.remove_connection(connection_id).await;
        tracing::info!("WebSocket 连接断开: {}", connection_id);
    }
    
    /// 处理 Hub 消息
    async fn handle_message(
        connection_manager: &ConnectionManagerHandle,
        connection_id: Uuid,
        message: HubMessage,
    ) -> anyhow::Result<()> {
        match message.message_type {
            MessageType::Invocation => {
                // 处理方法调用
                if let Some(target) = &message.target {
                    let arguments = message.arguments.unwrap_or_default();
                    
                    // 获取连接状态
                    let connection_state = connection_manager
                        .get_connection_state(connection_id)
                        .await
                        .ok_or_else(|| anyhow::anyhow!("连接状态不存在"))?;
                    
                    // 处理不同目标方法
                    match target.as_str() {
                        "TokenAuth" => {
                            Self::handle_token_auth(
                                connection_manager, 
                                connection_id, 
                                arguments
                            ).await?;
                        }
                        "ThingLoad" => {
                            Self::handle_thing_load(
                                connection_manager, 
                                connection_id, 
                                arguments,
                                connection_state.token
                            ).await?;
                        }
                        "CollectionLoad" => {
                            Self::handle_collection_load(
                                connection_manager, 
                                connection_id, 
                                arguments,
                                connection_state.token
                            ).await?;
                        }
                        _ => {
                            // 未知方法
                            let response = HubResponse::error(
                                message.invocation_id,
                                format!("未知方法: {}", target)
                            );
                            
                            connection_manager
                                .send_to_connection(connection_id, response)
                                .await?;
                        }
                    }
                }
            }
            MessageType::Ping => {
                // 响应 Ping 消息
                let response = HubResponse::success(message.invocation_id, serde_json::Value::Null);
                connection_manager
                    .send_to_connection(connection_id, response)
                    .await?;
            }
            _ => {
                tracing::warn!("未处理的消息类型: {:?}", message.message_type);
            }
        }
        
        Ok(())
    }
    
    /// 处理 Token 认证
    async fn handle_token_auth(
        connection_manager: &ConnectionManagerHandle,
        connection_id: Uuid,
        arguments: Vec<serde_json::Value>,
    ) -> anyhow::Result<()> {
        if arguments.is_empty() {
            return Err(anyhow::anyhow!("TokenAuth 缺少参数"));
        }
        
        // 解析 Token 参数
        let token_data = &arguments[0];
        
        // 从参数中提取 user_id 和 belong_id
        let user_id = token_data["user_id"]
            .as_u64()
            .ok_or_else(|| anyhow::anyhow!("无效的 user_id"))?;
            
        let belong_id = token_data["belong_id"]
            .as_u64()
            .ok_or_else(|| anyhow::anyhow!("无效的 belong_id"))?;
        
        // 创建 TokenModel
        let token_model = TokenModel {
            user_id,
            belong_id,
            // 其他字段根据实际需要添加
        };
        
        // 认证连接
        connection_manager
            .authenticate_connection(connection_id, token_model)
            .await?;
        
        // 发送认证成功响应（使用 "auth-1" 作为 invocation_id）
        let response = HubResponse::success(
            Some("auth-1".to_string()),
            serde_json::json!({
                "success": true,
                "message": "认证成功"
            })
        );
        
        connection_manager
            .send_to_connection(connection_id, response)
            .await?;
        
        Ok(())
    }
    
    /// 处理 Thing 加载
    async fn handle_thing_load(
        connection_manager: &ConnectionManagerHandle,
        connection_id: Uuid,
        arguments: Vec<serde_json::Value>,
        token: Option<TokenModel>,
    ) -> anyhow::Result<()> {
        // TODO: 实现 Thing 加载逻辑
        tracing::info!("处理 ThingLoad 请求: {:?}", arguments);
        
        // 模拟响应
        let response = HubResponse::success(
            None,
            serde_json::json!({
                "things": [],
                "total": 0
            })
        );
        
        connection_manager
            .send_to_connection(connection_id, response)
            .await?;
        
        Ok(())
    }
    
    /// 处理 Collection 加载
    async fn handle_collection_load(
        connection_manager: &ConnectionManagerHandle,
        connection_id: Uuid,
        arguments: Vec<serde_json::Value>,
        token: Option<TokenModel>,
    ) -> anyhow::Result<()> {
        // TODO: 实现 Collection 加载逻辑
        tracing::info!("处理 CollectionLoad 请求: {:?}", arguments);
        
        // 模拟响应
        let response = HubResponse::success(
            None,
            serde_json::json!({
                "collections": [],
                "total": 0
            })
        );
        
        connection_manager
            .send_to_connection(connection_id, response)
            .await?;
        
        Ok(())
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::services::{AuthConfig, AuthService, AuthServiceHandle, RealtimeAuthMiddleware};
    
    #[tokio::test]
    async fn test_websocket_handler_creation() {
        let manager = ConnectionManager::new(10);
        let handle = ConnectionManagerHandle::new(manager);
        
        // 创建认证中间件（需要AuthServiceHandle）
        let auth_config = AuthConfig::default();
        let auth_service = AuthService::new(auth_config).unwrap();
        let auth_service_handle = AuthServiceHandle::new(auth_service);
        let auth_middleware = RealtimeAuthMiddleware::new(auth_service_handle);
        
        let handler = WebSocketHandler::new(handle, auth_middleware);
        
        // 测试路由创建成功
        let route = handler.create_route();
        // 验证路由对象不为空（通过类型检查）
        assert!(std::mem::size_of_val(&route) > 0, "WebSocket 处理器路由对象应该被正确创建");
        
        // 验证路由处理器类型正确
        tracing::info!("WebSocket 处理器路由创建成功");
    }
}